import math
import numpy as np
from scipy.optimize import curve_fit


def sine(x, A, w, phi, c):
    return A * np.sin(w * x + phi) + c

# 输入绘制的坐标集合，返回拟合的坐标集合、倾角、曲线中心点坐标
def fit(data: np.array, r):
    # 提取x和y数据
    x_data = data[:, 0]  # 第一列作为x坐标
    y_data = data[:, 1]  # 第二列作为y坐标

    # 计算数据集的总行数
    n = len(data)

    # 找到中间索引
    middle_index = n // 2

    # 获取该索引对应的 x 和 y 坐标
    middle_point = data[middle_index]
    print(middle_point)

    maxA = max(y_data) - min(y_data)
    print("振幅应该为:", str(maxA))

    # 计算一个周期的宽度
    cycle_width = x_data[-1] - x_data[0]

    # 固定角频率 w 为一个周期对应的值
    w_fixed = 2*np.pi / cycle_width

    # 初始猜测值（A: 振幅, phi: 相位偏移, c: 垂直偏移）
    initial_guess = [(np.max(y_data) - np.min(y_data))/2, 0, np.mean(y_data)]

    # 使用 curve_fit 进行拟合，并增加最大迭代次数
    # 注意：因为 w 已经固定，所以我们不再需要它作为拟合参数传递
    params, params_covariance = curve_fit(lambda x, A, phi, c: sine(x, A, w_fixed, phi, c),
                                          x_data, y_data, p0=initial_guess, maxfev=10000)

    # 提取拟合参数
    A, phi, c = params

    # 如果振幅为负，将其变为正，并调整相位偏移
    if A < 0:
        A = -A
        phi += np.pi

    # 打印拟合结果
    print(f"拟合参数: 振幅(A)={A:.2f}, 角频率(w)={w_fixed:.8f}, 相位偏移(phi)={phi:.2f}, 垂直偏移(c)={c:.2f}")

    # 生成用于绘图的 x 轴数据
    x_fit = np.linspace(x_data.min(), x_data.max(), 1000)
    y_fit = sine(x_fit, A, w_fixed, phi, c)

    # 计算夹角
    angle = calculate_angle(A, w_fixed, phi, c)

    print("地层倾角为:", angle)

    return np.column_stack((x_fit, y_fit)),f"{angle:.1f}"+"度",middle_point


def calculate_angle(A, w, phi, c):
    # 计算最大斜率
    max_slope = A * w

    # 通过最大斜率计算夹角θ
    theta = math.atan(max_slope)

    # 将弧度转换为角度
    theta_degrees = math.degrees(theta)

    return theta_degrees


